import numpy as np
import matplotlib.pyplot as plt
import os

from vehicle_model import DoubleIntegratorVehicle
from lqr_controller import LQRController
from trajectory_generator import TrajectoryGenerator
from simulation import Simulation
from visualization import Visualization
from utils import set_plot_style

def run_test():
    """运行测试，比较不同控制参数的效果"""
    print("=" * 50)
    print("LQR轨迹跟踪测试")
    print("=" * 50)
    
    # 设置参数
    dt = 0.1
    sim_time = 30.0
    
    # 测试不同的Q和R组合
    param_sets = [
        {
            "name": "默认参数",
            "Q": np.diag([1000, 1000, 100, 100]),
            "R": np.diag([0.01, 0.01])
        },
        {
            "name": "超高位置权重",
            "Q": np.diag([5000, 5000, 500, 500]),
            "R": np.diag([0.001, 0.001])
        },
        {
            "name": "平衡权重",
            "Q": np.diag([500, 500, 50, 50]),
            "R": np.diag([0.1, 0.1])
        }
    ]
    
    # 存储所有结果
    all_results = []
    all_errors = []
    
    for params in param_sets:
        print("\n" + "=" * 30)
        print(f"测试参数组: {params['name']}")
        print("Q =\n", params["Q"])
        print("R =\n", params["R"])
        
        # 创建仿真环境
        sim = Simulation(dt=dt, sim_time=sim_time)
        
        # 设置控制器参数
        sim.controller.update_parameters(Q=params["Q"], R=params["R"])
        
        # 重置仿真
        sim.reset(trajectory_type='circle')
        
        # 运行仿真
        sim.run()
        
        # 获取结果
        results = sim.get_results()
        errors = sim.calculate_errors()
        
        # 存储结果
        all_results.append(results)
        all_errors.append(errors)
        
        print(f"RMSE: {errors['rmse']:.4f}")
        print(f"最大误差: {errors['max_error']:.4f}")
    
    # 绘制比较图
    plt.figure(figsize=(12, 8))
    
    # 绘制轨迹对比
    for i, (params, results) in enumerate(zip(param_sets, all_results)):
        ref_traj = results['reference_trajectory']
        actual_traj = results['states_history']
        
        plt.plot(ref_traj[:, 0], ref_traj[:, 1], 'k-', label='参考轨迹' if i == 0 else None)
        plt.plot(actual_traj[:, 0], actual_traj[:, 1], '--', 
                 label=f"{params['name']} (RMSE: {all_errors[i]['rmse']:.4f})")
    
    plt.xlabel('X位置')
    plt.ylabel('Y位置')
    plt.title('不同LQR参数的轨迹跟踪效果对比')
    plt.legend()
    plt.grid(True)
    plt.axis('equal')
    
    # 保存图像
    if not os.path.exists('results'):
        os.makedirs('results')
    plt.savefig('results/tracking_comparison.png', dpi=300, bbox_inches='tight')
    
    plt.show()
    
    # 返回最佳结果
    best_idx = np.argmin([e['rmse'] for e in all_errors])
    best_params = param_sets[best_idx]
    
    print("\n" + "=" * 50)
    print(f"最佳参数组: {best_params['name']}")
    print(f"RMSE: {all_errors[best_idx]['rmse']:.4f}")
    print("推荐命令:")
    print(f"python main.py --q-pos {best_params['Q'][0,0]} --q-vel {best_params['Q'][2,2]} --r-acc {best_params['R'][0,0]}")
    print("=" * 50)
    
    return best_params

if __name__ == "__main__":
    # 设置绘图样式
    set_plot_style()
    
    # 运行测试
    best_params = run_test() 